package main

import (
	"fmt"
)

/*
实现AVL树，实现思路：
1.定义最基本的数据结构
2.定义最基本的操作
3.用以上两个基本元素，将数据结构的各种操作描述出来
*/

/*
AVL树的实现，应该采用记录树高而非记录平衡因子的方式来实现，
因为二叉树的定义本身就是由树高之间的差不超过1来定义的，是
核心，而平衡因子只是树高的差，是表现；如果使用平衡因子来作
为节点数据成员，在旋转时，调整平衡因子就要处理非常多的分支，
算法实现非常繁杂，各种分支难以理解。
*/

// 树的节点定义
type AvlTreeNode struct {
	Data   interface{}
	Left   *AvlTreeNode
	Right  *AvlTreeNode
	Height int
}

// 树本身的定义
type AvlTree struct {
	// 树的根节点
	Root *AvlTreeNode
	// 比较函数，用于比较节点的值，初始化AVL树时提供
	CompareFunc func(interface{}, interface{}) int
}

// 计算节点的高度，入参可能为nil
func getHeight(node *AvlTreeNode) int {
	if node == nil {
		return 0
	} else {
		return node.Height
	}
}

// 获取v1和v2中的较大者
func max(v1, v2 int) int {
	if v1 > v2 {
		return v1
	} else {
		return v2
	}
}

// 返回AVL树中最大的节点
func getMaxNode(tree *AvlTree) *AvlTreeNode {
	if tree.Root == nil {
		return nil
	} else {
		maxNode := tree.Root
		for {
			if maxNode.Right != nil {
				maxNode = maxNode.Right
			} else {
				break
			}
		}
		return maxNode
	}
}

// 返回AVL树中最小的节点
func getMinNode(tree *AvlTree) *AvlTreeNode {
	if tree.Root == nil {
		return nil
	} else {
		minNode := tree.Root
		for {
			if minNode.Left != nil {
				minNode = minNode.Left
			} else {
				break
			}
		}
		return minNode
	}
}

// 右旋转，旋转后返回新的根节点
func (tree *AvlTree) RightRotate() {
	// 目前树的根节点
	cur_root := tree.Root

	// 调整节点的指针位置，进行右旋转
	new_root := cur_root.Left
	cur_root.Left = new_root.Right
	new_root.Right = cur_root

	// 节点的高度可能发生变化，先算右子树的高度，再算新的根节点的高度
	cur_root.Height = max(getHeight(cur_root.Left), getHeight(cur_root.Right)) + 1
	new_root.Height = max(getHeight(new_root.Left), cur_root.Height) + 1

	tree.Root = new_root
}

// 左旋转，旋转后返回新的根节点
func (tree *AvlTree) LeftRotate() {
	// 目前树的根节点
	cur_root := tree.Root
	// 调整节点的指针位置，进行左旋转
	new_root := cur_root.Right
	cur_root.Right = new_root.Left
	new_root.Left = cur_root

	// 节点的高度可能发生变化，先算左子树的高度，再算新的根节点的高度
	cur_root.Height = max(getHeight(cur_root.Left), getHeight(cur_root.Right)) + 1
	new_root.Height = max(getHeight(new_root.Right), cur_root.Height) + 1

	tree.Root = new_root
}

// 往二叉树插入节点
func (tree *AvlTree) Insert(data interface{}) {
	if tree.Root == nil {
		// 若树的根节点为空，插入值直接初始化根节点
		node := AvlTreeNode{
			Data:   data,
			Left:   nil,
			Right:  nil,
			Height: 1}
		tree.Root = &node
	} else {
		// 若树的根节点不为空，通过平衡二叉树的规则，对节点的值进行比较，决定插入左或右子树
		cmpRes := tree.CompareFunc(data, tree.Root.Data)
		if cmpRes < 0 {
			// 左子树进行节点插入
			leftRoot := &AvlTree{
				Root:        tree.Root.Left,
				CompareFunc: tree.CompareFunc}
			leftRoot.Insert(data)
			// 左子树的根节点可能已经发生了变化，需要更新原根节点的左孩子
			tree.Root.Left = leftRoot.Root
			// 若插入节点后，左子树的高度和右子树的高度相差2，AVL树需要进行旋转
			if getHeight(tree.Root.Left)-getHeight(tree.Root.Right) == 2 {
				// 被插入的节点在左子树的左子树，需要进行一次右旋转，恢复AVL树的特性
				if tree.CompareFunc(data, tree.Root.Left.Data) < 0 {
					tree.RightRotate()
				} else {
					// 被插入的节点在左子树的右子树，需要先进行一次左旋转，再进行一次右旋转
					tree.LeftRotate()
					tree.RightRotate()
				}
			}
		} else if cmpRes > 0 {
			// 右子树进行插入
			rightRoot := &AvlTree{
				Root:        tree.Root.Right,
				CompareFunc: tree.CompareFunc}
			rightRoot.Insert(data)
			// 右子树的根节点可能发生了变化，需要更新原根节点的右孩子
			tree.Root.Right = rightRoot.Root
			// 若插入节点后，右子树的高度和左子树的高度相差2，AVL树需要进行旋转
			if getHeight(tree.Root.Right)-getHeight(tree.Root.Left) == 2 {
				// 被插入的节点在右子树的右子树，需要进行一次左旋转，恢复AVL树的特性
				if tree.CompareFunc(data, tree.Root.Right.Data) > 0 {
					tree.LeftRotate()
				} else {
					// 被插入的节点在右子树的左子树，需要先进行一次右旋转，再进行一次左旋转
					tree.RightRotate()
					tree.LeftRotate()
				}
			}
		} else {
			// 此时要插入的节点和根节点相比值相同，不做任何处理
		}
		// 插入节点后，更新AVL树的高度
		tree.Root.Height = max(getHeight(tree.Root.Left), getHeight(tree.Root.Right)) + 1
	}

}

// 删除AVL树的节点
func (tree *AvlTree) Delete(data interface{}) {
	// 用二叉查找找到需要删除的节点
	if data != nil && tree.Root != nil {
		cmpRes := tree.CompareFunc(data, tree.Root.Data)
		if cmpRes < 0 {
			// 左子树进行节点删除
			leftRoot := &AvlTree{
				Root:        tree.Root.Left,
				CompareFunc: tree.CompareFunc}
			leftRoot.Delete(data)
			tree.Root.Left = leftRoot.Root
			// 删除后，右子树与左子树高度相差2，需要进行旋转，维持AVL树特性
			if getHeight(tree.Root.Right)-getHeight(tree.Root.Left) == 2 {
				if getHeight(tree.Root.Right.Left) > getHeight(tree.Root.Right.Right) {
					tree.RightRotate()
					tree.LeftRotate()
				} else {
					tree.LeftRotate()
				}
			}
		} else if cmpRes > 0 {
			// 右子树进行节点删除
			rightRoot := &AvlTree{
				Root:        tree.Root.Right,
				CompareFunc: tree.CompareFunc}
			rightRoot.Delete(data)
			tree.Root.Right = rightRoot.Root
			// 删除后，左子树与右子树高度相差2，需要进行旋转，维持AVL树特性
			if getHeight(tree.Root.Left)-getHeight(tree.Root.Right) == 2 {
				if getHeight(tree.Root.Left.Left) > getHeight(tree.Root.Left.Right) {
					tree.RightRotate()
				} else {
					tree.LeftRotate()
					tree.RightRotate()
				}
			}
		} else {
			// 当前就是要删除的节点
			if tree.Root.Left != nil && tree.Root.Right != nil {
				// 比较左右子树的高度
				if getHeight(tree.Root.Left) > getHeight(tree.Root.Right) {
					leftRoot := &AvlTree{Root: tree.Root.Left, CompareFunc: tree.CompareFunc}
					maxNode := getMaxNode(leftRoot)
					tree.Root.Data = maxNode.Data
					leftRoot.Delete(maxNode.Data)
					tree.Root.Left = leftRoot.Root
				} else {
					rightRoot := &AvlTree{Root: tree.Root.Right, CompareFunc: tree.CompareFunc}
					minNode := getMinNode(rightRoot)
					tree.Root.Data = minNode.Data
					rightRoot.Delete(minNode.Data)
					tree.Root.Right = rightRoot.Root
				}
			} else {
				// 左子树或右子树为空，
				tmpNode := tree.Root
				if tmpNode.Left != nil {
					tree.Root = tmpNode.Left
				} else if tmpNode.Right != nil {
					tree.Root = tmpNode.Right
				} else {
					tree.Root = nil
				}
			}
		}
	}
}

// 从AVL树中查找指定的节点
func (tree *AvlTree) Find(data interface{}) interface{} {
	node := tree.Root
	for node != nil {
		cmpRes := tree.CompareFunc(data, node.Data)
		if cmpRes < 0 {
			node = node.Left
		} else if cmpRes > 0 {
			node = node.Right
		} else {
			return node.Data
		}
	}
	return nil
}

// 将AvL树打印为字符串
func (tree *AvlTree) String() string {
	return ""
}

// 前序遍历
func (tree *AvlTree) PreOrderWalk(opFunc func(data interface{})) {
	if tree.Root != nil {
		opFunc(tree.Root.Data)
		leftRoot := &AvlTree{Root: tree.Root.Left, CompareFunc: tree.CompareFunc}
		leftRoot.PreOrderWalk(opFunc)
		rightRoot := &AvlTree{Root: tree.Root.Right, CompareFunc: tree.CompareFunc}
		rightRoot.PreOrderWalk(opFunc)
	} else {
		fmt.Print(" NULL ")
	}
}

// 初始化一颗空的AVL树
func New(compareFunc func(interface{}, interface{}) int) *AvlTree {
	return &AvlTree{
		Root:        nil,
		CompareFunc: compareFunc}
}

func main() {
	type kvData struct {
		key   int
		value int
	}

	compareFunc := func(d1, d2 interface{}) int {
		var kv1 kvData = d1.(kvData)
		var kv2 kvData = d2.(kvData)
		return kv1.key - kv2.key
	}
	tree := New(compareFunc)
	tree.Insert(kvData{1, 1})
	tree.Insert(kvData{2, 2})
	tree.Insert(kvData{3, 3})
	tree.Insert(kvData{4, 4})
	tree.Insert(kvData{5, 5})
	tree.Insert(kvData{6, 6})
	tree.Insert(kvData{7, 7})

	tree.Insert(kvData{15, 15})
	tree.Insert(kvData{16, 16})

	opFunc := func(data interface{}) {
		kvData := data.(kvData)
		fmt.Printf(" {key:%d, value:%d} ", kvData.key, kvData.value)
	}

	tree.PreOrderWalk(opFunc)
	fmt.Println()
	tree.Insert(kvData{14, 14})
	tree.PreOrderWalk(opFunc)
	fmt.Println()
	tree.Delete(kvData{key: 6})
	tree.Delete(kvData{key: 15})
	tree.PreOrderWalk(opFunc)
	fmt.Println()

}
